#include "fruit_detector/fruit_detector.hpp"
#include <memory>

// 包含生成的消息头文件
#include "fruit_detector/msg/fruit_info.hpp"
#include "fruit_detector/msg/detection_array.hpp"

namespace fruit_detector {

FruitDetector::FruitDetector(const rclcpp::NodeOptions &options)
    : Node("fruit_detector", options), 
      left_current_fps_(0.0f), right_current_fps_(0.0f),
      left_first_frame_(true), right_first_frame_(true),
      current_detect_part_(1) {
  // 声明参数
  this->declare_parameter("left_image_topic", "/left_cam/image_raw");
  this->declare_parameter("right_image_topic", "/right_cam/image_raw");
  this->declare_parameter("model_path_a", "");
  this->declare_parameter("model_path_b", "");
  this->declare_parameter("model_path_c", "");
  this->declare_parameter("conf_threshold", 0.25f);
  this->declare_parameter("iou_threshold", 0.45f);

  // 获取参数
  left_image_topic_ = this->get_parameter("left_image_topic").as_string();
  right_image_topic_ = this->get_parameter("right_image_topic").as_string();
  model_path_a_ = this->get_parameter("model_path_a").as_string();
  model_path_b_ = this->get_parameter("model_path_b").as_string();
  model_path_c_ = this->get_parameter("model_path_c").as_string();
  conf_threshold_ = this->get_parameter("conf_threshold").as_double();
  iou_threshold_ = this->get_parameter("iou_threshold").as_double();

  RCLCPP_INFO(this->get_logger(), "初始化检测器");
  RCLCPP_INFO(this->get_logger(), "模型A路径: %s", model_path_a_.c_str());
  RCLCPP_INFO(this->get_logger(), "模型B路径: %s", model_path_b_.c_str());
  RCLCPP_INFO(this->get_logger(), "模型C路径: %s", model_path_c_.c_str());

  // 定义三个模型的类别
  // 模型A：2个类别
  std::vector<std::string> classes_a = {"ripe", "unripe"};
  
  // 模型B：4个类别
  std::vector<std::string> classes_b = {
    "apple_ripe", "apple_unripe", "pear_ripe", "pear_unripe"
  };
  
  // 模型C：8个类别
  std::vector<std::string> classes_c = {
    "lajiao_ripe", "lajiao_unripe", "nangua_ripe", "nangua_unripe",
    "onion_ripe", "onion_unripe", "tomato_ripe", "tomato_unripe"
  };

  // 初始化左右相机的三个推理模型
  left_inference_a_ = std::make_unique<yolo::Inference>(model_path_a_, classes_a,
                                                        conf_threshold_, iou_threshold_);
  left_inference_b_ = std::make_unique<yolo::Inference>(model_path_b_, classes_b,
                                                        conf_threshold_, iou_threshold_);
  left_inference_c_ = std::make_unique<yolo::Inference>(model_path_c_, classes_c,
                                                        conf_threshold_, iou_threshold_);
  right_inference_a_ = std::make_unique<yolo::Inference>(model_path_a_, classes_a,
                                                         conf_threshold_, iou_threshold_);
  right_inference_b_ = std::make_unique<yolo::Inference>(model_path_b_, classes_b,
                                                         conf_threshold_, iou_threshold_);
  right_inference_c_ = std::make_unique<yolo::Inference>(model_path_c_, classes_c,
                                                         conf_threshold_, iou_threshold_);

  // 创建左相机发布者
  left_image_publisher_ = this->create_publisher<sensor_msgs::msg::Image>(
      "/left_detection_image", 10);
  left_detection_publisher_ = this->create_publisher<fruit_detector::msg::DetectionArray>(
      "/detect_left", 10);

  // 创建右相机发布者
  right_image_publisher_ = this->create_publisher<sensor_msgs::msg::Image>(
      "/right_detection_image", 10);
  right_detection_publisher_ = this->create_publisher<fruit_detector::msg::DetectionArray>(
      "/detect_right", 10);

  // 创建订阅者
  left_image_subscription_ = this->create_subscription<sensor_msgs::msg::Image>(
      left_image_topic_, 10,
      std::bind(&FruitDetector::leftImageCallback, this, std::placeholders::_1));

  right_image_subscription_ = this->create_subscription<sensor_msgs::msg::Image>(
      right_image_topic_, 10,
      std::bind(&FruitDetector::rightImageCallback, this, std::placeholders::_1));

  // 创建detect_part订阅者
  detect_part_subscription_ = this->create_subscription<std_msgs::msg::Int32>(
      "/detect_part", 10,
      std::bind(&FruitDetector::detectPartCallback, this, std::placeholders::_1));

  RCLCPP_INFO(this->get_logger(), "水果检测节点初始化完成");
  RCLCPP_INFO(this->get_logger(), "订阅左相机: %s", left_image_topic_.c_str());
  RCLCPP_INFO(this->get_logger(), "订阅右相机: %s", right_image_topic_.c_str());
  RCLCPP_INFO(this->get_logger(), "订阅detect_part话题: /detect_part");
  RCLCPP_INFO(this->get_logger(), "默认使用模型A (detect_part=1)");
}

// detect_part话题回调
void FruitDetector::detectPartCallback(const std_msgs::msg::Int32::SharedPtr msg) {
  int part = msg->data;
  if (part >= 1 && part <= 3) {
    current_detect_part_ = part;
    const char* model_name[] = {"", "A", "B", "C"};
    RCLCPP_INFO(this->get_logger(), "切换到模型%s (detect_part=%d)", 
                model_name[part], part);
  } else {
    RCLCPP_WARN(this->get_logger(), "无效的detect_part值: %d (有效值为1, 2, 3)", part);
  }
}

// 左相机图像回调
void FruitDetector::leftImageCallback(const sensor_msgs::msg::Image::SharedPtr msg) {
  // 根据current_detect_part_选择对应的推理实例
  std::unique_ptr<yolo::Inference>* inference_ptr = nullptr;
  switch (current_detect_part_) {
    case 1:
      inference_ptr = &left_inference_a_;
      break;
    case 2:
      inference_ptr = &left_inference_b_;
      break;
    case 3:
      inference_ptr = &left_inference_c_;
      break;
    default:
      RCLCPP_ERROR(this->get_logger(), "无效的detect_part值: %d", current_detect_part_);
      return;
  }
  processImage(msg, *inference_ptr, left_image_publisher_, left_detection_publisher_);
}

// 右相机图像回调
void FruitDetector::rightImageCallback(const sensor_msgs::msg::Image::SharedPtr msg) {
  // 根据current_detect_part_选择对应的推理实例
  std::unique_ptr<yolo::Inference>* inference_ptr = nullptr;
  switch (current_detect_part_) {
    case 1:
      inference_ptr = &right_inference_a_;
      break;
    case 2:
      inference_ptr = &right_inference_b_;
      break;
    case 3:
      inference_ptr = &right_inference_c_;
      break;
    default:
      RCLCPP_ERROR(this->get_logger(), "无效的detect_part值: %d", current_detect_part_);
      return;
  }
  processImage(msg, *inference_ptr, right_image_publisher_, right_detection_publisher_);
}

// 处理图像并进行推理
void FruitDetector::processImage(
    const sensor_msgs::msg::Image::SharedPtr msg,
    const std::unique_ptr<yolo::Inference>& inference,
    const rclcpp::Publisher<sensor_msgs::msg::Image>::SharedPtr& image_publisher,
    const rclcpp::Publisher<DetectionArrayMsg>::SharedPtr& detection_publisher) {
  try {
    // 判断是左相机还是右相机
    bool is_left_camera = (image_publisher == left_image_publisher_);
    
    // 计算帧率
    auto now = std::chrono::steady_clock::now();
    float& current_fps = is_left_camera ? left_current_fps_ : right_current_fps_;
    bool& first_frame = is_left_camera ? left_first_frame_ : right_first_frame_;
    auto& last_frame_time = is_left_camera ? last_left_frame_time_ : last_right_frame_time_;
    
    if (!first_frame) {
      auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
          now - last_frame_time).count();
      if (elapsed > 0) {
        current_fps = 1000.0f / static_cast<float>(elapsed);
      }
    } else {
      first_frame = false;
    }
    last_frame_time = now;
    
    // 转换ROS图像消息为OpenCV格式
    cv_bridge::CvImagePtr cv_ptr =
        cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
    cv::Mat frame = cv_ptr->image;

    // 获取图像尺寸
    int img_width = frame.cols;
    int img_height = frame.rows;
    int center_x = img_width / 2;
    int center_y = img_height / 2;

    // 执行推理
    inference->RunInference(frame);

    // 获取检测结果
    auto detections = inference->GetContours();

    // 绘制中心参考线
    cv::line(frame, cv::Point(center_x, 0), cv::Point(center_x, img_height), 
             cv::Scalar(0, 255, 0), 1);
    cv::line(frame, cv::Point(0, center_y), cv::Point(img_width, center_y), 
             cv::Scalar(0, 255, 0), 1);
    
    // 处理检测结果
    for (auto& detection : detections) {
      if (detection.empty()) continue;
      
      int class_id = static_cast<int>(detection[0].x);
      float confidence = detection[0].y;
      
      // 计算中心点
      float center_x = (detection[1].x + detection[2].x) / 2;
      float center_y = (detection[1].y + detection[2].y) / 2;
      
      // 在图像上绘制中心点
      cv::circle(frame, cv::Point(center_x, center_y), 5,
                 cv::Scalar(0, 0, 255), -1);
    }
    
    // 添加FPS显示
    char fps_text[50];
    snprintf(fps_text, sizeof(fps_text), "FPS: %.1f", current_fps);
    cv::putText(frame, fps_text, cv::Point(20, 50), cv::FONT_HERSHEY_SIMPLEX, 
                1.0, cv::Scalar(0, 255, 0), 2);
    
    // 添加相机标识
    std::string camera_text = is_left_camera ? "Left Camera" : "Right Camera";
    cv::putText(frame, camera_text, cv::Point(20, 100), cv::FONT_HERSHEY_SIMPLEX,
                1.0, cv::Scalar(0, 0, 255), 2);

    // 创建并发布水果信息消息
    auto fruit_info_msg = createFruitInfoMessage(detections, msg->header, img_width, img_height);
    detection_publisher->publish(fruit_info_msg);

    // 发布带有检测结果的图像
    cv_bridge::CvImage detection_image;
    detection_image.header = msg->header;
    detection_image.encoding = sensor_msgs::image_encodings::BGR8;
    detection_image.image = frame;

    // 在发布前确认图像有效
    if (!frame.empty()) {
      auto image_msg = detection_image.toImageMsg();
      image_publisher->publish(*image_msg);
    } else {
      RCLCPP_WARN(this->get_logger(), "检测结果图像为空，无法发布");
    }

  } catch (cv_bridge::Exception &e) {
    RCLCPP_ERROR(this->get_logger(), "CV桥接异常: %s", e.what());
  } catch (std::exception &e) {
    RCLCPP_ERROR(this->get_logger(), "处理异常: %s", e.what());
  }
}

// 创建水果信息消息
FruitDetector::DetectionArrayMsg FruitDetector::createFruitInfoMessage(
    const std::vector<std::vector<cv::Point2f>> &detections,
    const std_msgs::msg::Header &header,
    int img_width,
    int img_height) {
  DetectionArrayMsg detection_array;
  
  for (const auto &detection : detections) {
    if (detection.empty()) {
      continue; // 跳过空的检测
    }

    fruit_detector::msg::FruitInfo fruit_info;
    
    // 提取类别和置信度 (检测器存储在contour的第一个点)
    int class_id = static_cast<int>(detection[0].x);
    float confidence = detection[0].y;
    
    // 计算中心点和边界框尺寸
    float center_x = (detection[1].x + detection[2].x) / 2;
    float center_y = (detection[1].y + detection[2].y) / 2;
    
    // 计算偏移量（相对于图像中心）
    float img_center_x = img_width / 2;
    float img_center_y = img_height / 2;
    float x_offset = center_x - img_center_x;
    float y_offset = center_y - img_center_y;

    // 设置水果信息
    fruit_info.type = class_id;
    fruit_info.offset_x = x_offset;
    fruit_info.offset_y = y_offset;
    fruit_info.isgood = true;  // 默认设置为好水果
    
    detection_array.fruits.push_back(fruit_info);
  }

  return detection_array;
}

} // namespace fruit_detector

#include "rclcpp_components/register_node_macro.hpp"
RCLCPP_COMPONENTS_REGISTER_NODE(fruit_detector::FruitDetector)